from torchviz import make_dot
import torch
from EfficientNetV2WithGCA import EfficientNetV2WithGCA
import graphviz


model = EfficientNetV2WithGCA()
x = torch.randn(1, 3, 224, 224)  # 输入示例
y = model(x)

# 生成流程图
dot = make_dot(y, params=dict(model.named_parameters()))
dot.render("EfficientNetV2_GCA", format="png")  # 保存为 PNG 图像
